﻿#include "uaesencrypt.h"

#include <openssl/evp.h>
#include <openssl/err.h>

#include <QString>
#include <QDebug>

struct ScopedPointerCipherCtxDeleter
{
    static inline void cleanup(EVP_CIPHER_CTX *pointer)
    {
        EVP_CIPHER_CTX_free(pointer);
    }
};

class UAESEncryptPrivate
{
    Q_DECLARE_PUBLIC(UAESEncrypt)
public:
    UAESEncryptPrivate(UAESEncrypt* parent) : q_ptr(parent) {}
    ~UAESEncryptPrivate() {}

    bool isVaild(const QByteArray &key, const QByteArray &iv);

    const EVP_CIPHER *evpcipher();

    const EVP_CIPHER *ecbevpcipher();

    const EVP_CIPHER *cbcevpcipher();

    const EVP_CIPHER *cfbevpcipher();

    const EVP_CIPHER *ofbevpcipher();

    int padding();

private:
    UAESEncrypt* q_ptr = nullptr;
    UAESEncrypt::AESType aesType = UAESEncrypt::AESType::AES_256;
    UAESEncrypt::AESMode aesMode = UAESEncrypt::AESMode::AES_CBC;
    UAESEncrypt::AESPadding aesPadding = UAESEncrypt::AESPadding::AES_PKCS7;
};

bool UAESEncryptPrivate::isVaild(const QByteArray &key, const QByteArray &iv)
{
    bool iskeyVaild = true;
    switch (aesType) {
    case UAESEncrypt::AES_128:
        iskeyVaild = (key.size() == 16);
        break;
    case UAESEncrypt::AES_192:
        iskeyVaild = (key.size() == 24);
        break;
    case UAESEncrypt::AES_256:
        iskeyVaild = (key.size() == 32);
        break;
    default:
        break;
    }
    return iskeyVaild && (iv.size() == 16);
}

const EVP_CIPHER *UAESEncryptPrivate::evpcipher()
{
    const EVP_CIPHER *cipher = nullptr;
    switch (aesMode) {
    case UAESEncrypt::AES_ECB:
        cipher = ecbevpcipher();
        break;
    case UAESEncrypt::AES_CBC:
        cipher = cbcevpcipher();
        break;
    case UAESEncrypt::AES_CFB:
        cipher = cfbevpcipher();
        break;
    case UAESEncrypt::AES_OFB:
        cipher = ofbevpcipher();
        break;
    default:
        break;
    }
    return cipher;
}

const EVP_CIPHER *UAESEncryptPrivate::ecbevpcipher()
{
    const EVP_CIPHER *cipher = nullptr;
    switch (aesType) {
    case UAESEncrypt::AES_128:
        cipher = EVP_aes_128_ecb();
        break;
    case UAESEncrypt::AES_192:
        cipher = EVP_aes_192_ecb();
        break;
    case UAESEncrypt::AES_256:
        cipher = EVP_aes_256_ecb();
        break;
    default:
        break;
    }
    return cipher;
}

const EVP_CIPHER *UAESEncryptPrivate::cbcevpcipher()
{
    const EVP_CIPHER *cipher = nullptr;
    switch (aesType) {
    case UAESEncrypt::AES_128:
        cipher = EVP_aes_128_cbc();
        break;
    case UAESEncrypt::AES_192:
        cipher = EVP_aes_192_cbc();
        break;
    case UAESEncrypt::AES_256:
        cipher = EVP_aes_256_cbc();
        break;
    default:
        break;
    }
    return cipher;
}

const EVP_CIPHER *UAESEncryptPrivate::cfbevpcipher()
{
    const EVP_CIPHER *cipher = nullptr;
    switch (aesType) {
    case UAESEncrypt::AES_128:
        cipher = EVP_aes_128_cfb();
        break;
    case UAESEncrypt::AES_192:
        cipher = EVP_aes_192_cfb();
        break;
    case UAESEncrypt::AES_256:
        cipher = EVP_aes_256_cfb();
        break;
    default:
        break;
    }
    return cipher;
}

const EVP_CIPHER *UAESEncryptPrivate::ofbevpcipher()
{
    const EVP_CIPHER *cipher = nullptr;
    switch (aesType) {
    case UAESEncrypt::AES_128:
        cipher = EVP_aes_128_ofb();
        break;
    case UAESEncrypt::AES_192:
        cipher = EVP_aes_192_ofb();
        break;
    case UAESEncrypt::AES_256:
        cipher = EVP_aes_256_ofb();
        break;
    default:
        break;
    }
    return cipher;
}

int UAESEncryptPrivate::padding()
{
    int pad = -1;
    switch (aesPadding) {
    case UAESEncrypt::AES_PKCS7:
        pad = EVP_PADDING_PKCS7;
        break;
    case UAESEncrypt::AES_ISO7816_4:
        pad = EVP_PADDING_ISO7816_4;
        break;
    case UAESEncrypt::AES_ANSI923:
        pad = EVP_PADDING_ANSI923;
        break;
    case UAESEncrypt::AES_ISO10126:
        pad = EVP_PADDING_ISO10126;
        break;
    case UAESEncrypt::AES_ZERO:
        pad = EVP_PADDING_ZERO;
        break;
    default:
        break;
    }
    return pad;
}

UAESEncrypt::UAESEncrypt() : d_ptr(new UAESEncryptPrivate(this))
{
}

UAESEncrypt::UAESEncrypt(AESType type, AESMode mode, AESPadding padding)
    : d_ptr(new UAESEncryptPrivate(this))
{
    Q_D(UAESEncrypt);
    d->aesType = type;
    d->aesMode = mode;
    d->aesPadding = padding;
}

UAESEncrypt::~UAESEncrypt()
{
}

QByteArray UAESEncrypt::decrypt(QByteArray cyphertext, const QByteArray &key, const QByteArray &iv)
{
    Q_D(UAESEncrypt);
    ERR_clear_error();
    if (!d->isVaild(key, iv))
    {
        return QByteArray();
    }

    QByteArray plaintext;
    QScopedPointer<EVP_CIPHER_CTX, ScopedPointerCipherCtxDeleter> ctx(EVP_CIPHER_CTX_new());

    if (!EVP_DecryptInit_ex(ctx.data(), d->evpcipher(), nullptr, (unsigned char *)key.data(), (unsigned char *)iv.data()))
    {
        plaintext.clear();
        return plaintext;
    }

    if (!EVP_CIPHER_CTX_set_padding(ctx.data(), d->padding()))
    {
        plaintext.clear();
        return plaintext;
    }

    int outlen = 0;
    plaintext.resize(cyphertext.size());
    if (!EVP_DecryptUpdate(ctx.data(), (unsigned char *)plaintext.data(), &outlen, (unsigned char *)cyphertext.data(), cyphertext.size()))
    {
        plaintext.clear();
        return plaintext;
    }
    int tmplen = 0;
    if (!EVP_DecryptFinal_ex(ctx.data(), (unsigned char *)plaintext.data() + outlen, &tmplen))
    {
        plaintext.clear();
        return plaintext;
    }
    return plaintext.left(outlen + tmplen);
}

QByteArray UAESEncrypt::encrypt(QByteArray plaintext, const QByteArray &key, const QByteArray &iv)
{
    Q_D(UAESEncrypt);
    ERR_clear_error();
    if (!d->isVaild(key, iv))
    {
        return QByteArray();
    }

    QByteArray cyphertext;
    QScopedPointer<EVP_CIPHER_CTX, ScopedPointerCipherCtxDeleter> ctx(EVP_CIPHER_CTX_new());

    if (!EVP_EncryptInit_ex(ctx.data(), d->evpcipher(), nullptr, (unsigned char*)key.data(), (unsigned char*)iv.data()))
    {
        cyphertext.clear();
        return cyphertext;
    }

    if (!EVP_CIPHER_CTX_set_padding(ctx.data(), d->padding()))
    {
        cyphertext.clear();
        return cyphertext;
    }

    int outlen = 0;
    cyphertext.resize(plaintext.size() + 1024);
    if (!EVP_EncryptUpdate(ctx.data(), (unsigned char *)cyphertext.data(), &outlen, (const unsigned char *)plaintext.data(), plaintext.size()))
    {
        cyphertext.clear();
        return cyphertext;
    }
    int tmplen = 0;
    if (!EVP_EncryptFinal_ex(ctx.data(), (unsigned char *)cyphertext.data() + outlen, &tmplen))
    {
        cyphertext.clear();
        return cyphertext;
    }

    return cyphertext.left(outlen + tmplen);
}

QString UAESEncrypt::GetLastError()
{
    unsigned long lastError = ERR_peek_last_error();
    if (lastError == 0) {
        return "";
    }
    char errorString[256];
    ERR_error_string_n(lastError, errorString, sizeof(errorString));
    return errorString;
}
